"""
# k-means算法原理
1. 核心思想：将n个样本划分为k个簇，使簇内样本相似度高（距离近），簇间相似度低（距离远）
2. 核心步骤（迭代优化）：
   - 步骤1：初始化k个聚类中心（随机选择k个样本作为初始中心）
   - 步骤2：分配样本：计算每个样本到所有中心的欧氏距离，将样本归入距离最近的簇
   - 步骤3：更新中心：计算每个簇内所有样本的特征均值，作为新的聚类中心
   - 步骤4：重复步骤2-3，直到中心不再变化或达到最大迭代次数
3. 关键指标：
   - 欧氏距离：衡量样本与中心的相似度，公式：d(x,y)=√[(x1-y1)²+(x2-y2)²]
   - 轮廓系数：评估聚类效果，取值[-1,1]，越接近1说明聚类效果越好（纯numpy实现）
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 设置中文字体（避免可视化乱码）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# ---------------------- 1. 生成聚类数据集 ----------------------
def generate_kmeans_dataset(num_samples=500, num_clusters=3, noise=0.8):
    """生成带真实标签的聚类数据集并保存为CSV"""
    # 定义3个聚类中心（二维特征）
    centers = np.array([[2, 3], [7, 5], [1, 8]])
    # 每个簇的样本数（尽量平均分配）
    samples_per_cluster = num_samples // num_clusters
    remaining = num_samples % num_clusters
    
    X = []
    true_labels = []
    for i in range(num_clusters):
        # 生成簇内样本（高斯分布，添加噪声）
        num = samples_per_cluster + (1 if i < remaining else 0)
        cluster_X = np.random.normal(centers[i], noise, (num, 2))
        X.append(cluster_X)
        true_labels.extend([i]*num)
    
    # 合并数据并打乱
    X = np.vstack(X)
    true_labels = np.array(true_labels)
    shuffle_idx = np.random.permutation(num_samples)
    X_shuffled = X[shuffle_idx]
    true_labels_shuffled = true_labels[shuffle_idx]
    
    # 保存为CSV
    df = pd.DataFrame({
        'x1': X_shuffled[:, 0],
        'x2': X_shuffled[:, 1],
        'true_label': true_labels_shuffled
    })
    df.to_csv('kmeans_dataset.csv', index=False)
    print(f"k-means数据集生成完成：{num_samples}个样本，{num_clusters}个真实簇，保存为kmeans_dataset.csv")
    
    # 可视化真实数据分布
    plt.figure(figsize=(8, 6))
    for i in range(num_clusters):
        mask = true_labels_shuffled == i
        plt.scatter(df.loc[mask, 'x1'], df.loc[mask, 'x2'], label=f'真实簇{i+1}', alpha=0.7)
    plt.scatter(centers[:, 0], centers[:, 1], c='black', marker='*', s=200, label='真实中心')
    plt.xlabel('特征x1')
    plt.ylabel('特征x2')
    plt.title('k-means数据集真实分布')
    plt.legend()
    plt.savefig('kmeans_true_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("真实数据分布图已保存为kmeans_true_distribution.png")
    
    return X_shuffled, true_labels_shuffled

# 生成数据集（若已生成，可注释此句直接加载）
X, true_labels = generate_kmeans_dataset(num_samples=500, num_clusters=3)
# 加载数据集（若已生成，用此句替换上一句）
# df = pd.read_csv('kmeans_dataset.csv')
# X = df[['x1', 'x2']].values
# true_labels = df['true_label'].values

print(f"加载数据：{X.shape[0]}个样本，{X.shape[1]}个特征")

# ---------------------- 2. 核心工具函数（无第三方库） ----------------------
def euclidean_distance(x1, x2):
    """计算两个样本的欧氏距离"""
    return np.sqrt(np.sum((x1 - x2)**2))

def calculate_silhouette_score(X, labels):
    """纯numpy实现轮廓系数计算（无sklearn依赖）
    轮廓系数s(i) = (b(i) - a(i)) / max(a(i), b(i))
    - a(i)：样本i到同一簇内所有其他样本的平均距离（簇内相似度）
    - b(i)：样本i到所有其他簇的平均距离的最小值（簇间相似度）
    - 整体轮廓系数：所有样本s(i)的平均值
    """
    num_samples = len(X)
    silhouette_scores = []
    
    for i in range(num_samples):
        xi = X[i]
        cluster_i = labels[i]
        
        # 1. 计算a(i)：簇内平均距离
        same_cluster_samples = X[labels == cluster_i]  # 同一簇的所有样本
        if len(same_cluster_samples) == 1:
            # 若簇内只有1个样本，a(i)=0（无其他样本可计算距离）
            a_i = 0.0
        else:
            # 计算xi到同一簇其他样本的距离均值
            distances_same = [euclidean_distance(xi, xj) for xj in same_cluster_samples if not np.all(xi == xj)]
            a_i = np.mean(distances_same)
        
        # 2. 计算b(i)：簇间最小平均距离
        other_clusters = np.unique(labels[labels != cluster_i])  # 所有其他簇的标签
        if len(other_clusters) == 0:
            # 若只有1个簇，b(i)=0（无其他簇可比较）
            b_i = 0.0
        else:
            avg_distances_other = []
            for cluster_j in other_clusters:
                # 计算xi到簇j所有样本的平均距离
                other_cluster_samples = X[labels == cluster_j]
                distances_other = [euclidean_distance(xi, xj) for xj in other_cluster_samples]
                avg_distances_other.append(np.mean(distances_other))
            # 取最小的簇间平均距离作为b(i)
            b_i = np.min(avg_distances_other)
        
        # 3. 计算样本i的轮廓系数（避免除以0）
        if max(a_i, b_i) == 0:
            s_i = 0.0
        else:
            s_i = (b_i - a_i) / max(a_i, b_i)
        silhouette_scores.append(s_i)
    
    # 4. 整体轮廓系数（所有样本的均值）
    return np.mean(silhouette_scores)

# ---------------------- 3. k-means核心实现 ----------------------
class KMeansClustering:
    def __init__(self, n_clusters=3, max_iter=100, tol=1e-4):
        self.n_clusters = n_clusters  # 簇数k
        self.max_iter = max_iter      # 最大迭代次数
        self.tol = tol                # 中心变化阈值（小于此值则停止迭代）
        self.centers = None           # 最终聚类中心
        self.labels = None            # 样本的聚类标签

    def initialize_centers(self, X):
        """随机初始化k个聚类中心（从样本中选择）"""
        idx = np.random.choice(X.shape[0], self.n_clusters, replace=False)
        return X[idx]

    def assign_labels(self, X, centers):
        """为每个样本分配簇标签（距离最近的中心）"""
        labels = []
        for x in X:
            # 计算到所有中心的距离
            distances = [euclidean_distance(x, center) for center in centers]
            # 选择距离最小的中心的索引作为标签
            labels.append(np.argmin(distances))
        return np.array(labels)

    def update_centers(self, X, labels):
        """更新聚类中心（簇内样本的均值）"""
        centers = []
        for i in range(self.n_clusters):
            # 取该簇所有样本
            cluster_samples = X[labels == i]
            # 计算均值（若簇为空，保留原中心）
            if len(cluster_samples) == 0:
                centers.append(self.centers[i])
            else:
                centers.append(np.mean(cluster_samples, axis=0))
        return np.array(centers)

    def fit(self, X):
        """训练k-means模型"""
        # 步骤1：初始化中心
        self.centers = self.initialize_centers(X)
        print(f"初始聚类中心：\n{self.centers}")

        # 步骤2：迭代分配标签和更新中心
        for iter in range(self.max_iter):
            # 保存当前中心（用于判断是否收敛）
            old_centers = self.centers.copy()
            
            # 分配标签
            self.labels = self.assign_labels(X, self.centers)
            
            # 更新中心
            self.centers = self.update_centers(X, self.labels)
            
            # 计算中心变化量（所有中心的欧氏距离之和）
            center_change = np.sum([euclidean_distance(old, new) for old, new in zip(old_centers, self.centers)])
            
            # 收敛判断：若中心变化量小于阈值，停止迭代
            if center_change < self.tol:
                print(f"迭代{iter+1}次后收敛（中心变化量={center_change:.6f} < 阈值{self.tol}）")
                break
        else:
            print(f"达到最大迭代次数{self.max_iter}，停止迭代")

        return self.labels, self.centers

# ---------------------- 4. 实验结果与分析 ----------------------
# 超参数设置（与真实簇数一致，实际应用中需通过肘部法则选择k）
N_CLUSTERS = 3
MAX_ITER = 100
TOL = 1e-4

# 训练k-means模型
kmeans = KMeansClustering(n_clusters=N_CLUSTERS, max_iter=MAX_ITER, tol=TOL)
pred_labels, final_centers = kmeans.fit(X)

# 计算聚类效果指标（纯numpy实现的轮廓系数）
sil_score = calculate_silhouette_score(X, pred_labels)

# 可视化聚类结果
plt.figure(figsize=(8, 6))
for i in range(N_CLUSTERS):
    mask = pred_labels == i
    plt.scatter(X[mask, 0], X[mask, 1], label=f'预测簇{i+1}', alpha=0.7)
# 绘制最终聚类中心
plt.scatter(final_centers[:, 0], final_centers[:, 1], c='red', marker='*', s=200, label='预测中心')
plt.xlabel('特征x1')
plt.ylabel('特征x2')
plt.title(f'k-means聚类结果（k={N_CLUSTERS}，轮廓系数={sil_score:.4f}）')
plt.legend()
plt.savefig('kmeans_clustering_result.png', dpi=300, bbox_inches='tight')
plt.close()

# 打印结果
print(f"\n=== k-means实验结果 ===")
print(f"1. 最终聚类中心：")
for i, center in enumerate(final_centers, start=1):
    print(f"   簇{i}中心：({center[0]:.4f}, {center[1]:.4f})")

print(f"\n2. 聚类效果评估：")
print(f"   轮廓系数：{sil_score:.4f}（越接近1，聚类效果越好）")

# 统计每个簇的样本数
cluster_counts = np.bincount(pred_labels)
print(f"\n3. 各簇样本数：")
for i, count in enumerate(cluster_counts, start=1):
    print(f"   簇{i}：{count}个样本（占比{count/X.shape[0]:.2%}）")

# 结果分析：
# - 轮廓系数≈0.6-0.8，说明聚类效果良好（因数据带噪声，与真实簇有少量重叠）
# - 预测中心与真实中心（[2,3],[7,5],[1,8]）接近，证明算法能有效找到簇中心
# - 实际应用中，若未知k，可通过“肘部法则”（绘制k与簇内平方和的曲线）选择最优k